# import websocket
# import json
# import os
# import hashlib
# import uuid
# import base64
#
# WS_IP_ADDRESS_PORT = "ws://47.106.145.145:50232"
# #WS_IP_ADDRESS_PORT = "ws://localhost:50232"
#
# def calculate_md5(file_path):
#     md5_hash = hashlib.md5()
#     with open(file_path, "rb") as f:
#         for chunk in iter(lambda: f.read(4096), b""):
#             md5_hash.update(chunk)
#     return md5_hash.hexdigest()
#
# def on_message(ws, message):
#     """处理服务器发送的消息（包括回传文件）"""
#     try:
#         frame = json.loads(message)
#         header = frame.get("header", {})
#         payload = frame.get("payload", {})
#         status = header.get("status")
#
#         # 处理服务器回传的文件
#         if status == 3:
#             # 元数据帧：初始化接收
#             global current_file
#             current_file = {
#                 "uuid": header["uuid"],
#                 "filename": header["fileName"],
#                 "total_size": int(header["size"]),
#                 "md5": header["MD5"],
#                 "total_chunks": payload["packageNum"],
#                 "chunks": {}
#             }
#             print(f"开始接收服务器回传文件: {current_file['filename']}")
#
#         elif status == 4:
#             # 数据帧：收集数据块
#             if not current_file:
#                 print("未收到元数据帧，忽略数据块")
#                 return
#
#             chunk_num = payload["num"]
#             data = base64.b64decode(payload["data"])
#             current_file["chunks"][chunk_num] = data
#             print(f"收到回传数据块 {chunk_num}/{current_file['total_chunks']}", end='\r')
#
#         elif status == 5:
#             # 结束帧：合并文件并校验
#             if not current_file:
#                 print("未收到元数据帧，忽略结束帧")
#                 return
#
#             # 合并文件
#             save_dir = "./received_from_server"
#             os.makedirs(save_dir, exist_ok=True)
#             file_path = os.path.join(save_dir, current_file["filename"])
#
#             with open(file_path, "wb") as f:
#                 for i in range(current_file["total_chunks"]):
#                     if i not in current_file["chunks"]:
#                         print(f"缺失数据块 {i}，文件不完整")
#                         return
#                     f.write(current_file["chunks"][i])
#
#             # 校验MD5
#             actual_md5 = calculate_md5(file_path)
#             if actual_md5 == current_file["md5"]:
#                 print(f"\n文件接收完成并校验通过: {file_path}")
#             else:
#                 print(f"\nMD5校验失败！预期: {current_file['md5']}, 实际: {actual_md5}")
#                 os.remove(file_path)
#
#             # 重置当前文件状态
#             #global current_file
#             current_file = None
#
#         elif status == 99:
#             print("===>"+str(frame))
#
#     except Exception as e:
#         print(f"处理消息出错: {e}")
#
# def send_file_with_protocol(file_path, server_url=WS_IP_ADDRESS_PORT):
#     if not os.path.isfile(file_path):
#         print(f"文件不存在: {file_path}")
#         return
#
#     # 计算文件元数据
#     file_size = os.path.getsize(file_path)
#     file_md5 = calculate_md5(file_path)
#     file_uuid = str(uuid.uuid4())  # 生成唯一UUID
#     chunk_size = 1024*32
#     total_chunks = (file_size + chunk_size - 1) // chunk_size  # 总块数
#
#     print(f"准备发送文件: {file_path}")
#     print(f"大小: {file_size} bytes, MD5: {file_md5}, UUID: {file_uuid}")
#
#     def on_open(ws):
#         # 1. 发送第一帧（元数据）
#         first_frame = {
#             "header": {
#                 "user_id": "default_user",
#                 "uuid": file_uuid,
#                 "msgType": "txt",
#                 "fileName": file_path.rsplit('/', 1)[-1],
#                 "size": str(file_size),
#                 "MD5": file_md5,
#                 "status": 0
#             },
#             "payload": {
#                 "packageNum": total_chunks,
#                 "data": ""
#             }
#         }
#         ws.send(json.dumps(first_frame))
#         print("已发送第一帧（元数据）")
#         print(first_frame)
#
#         # 2. 发送数据帧
#         with open(file_path, "rb") as f:
#             for chunk_num in range(total_chunks):
#                 # 读取8192字节并Base64编码
#                 data = f.read(chunk_size)
#                 b64_data = base64.b64encode(data).decode('utf-8').replace('\n', '').replace('\r', '')
#
#                 data_frame = {
#                     "header": {
#                         "user_id": "default_user",
#                         "uuid": file_uuid,
#                         "status": 1
#                     },
#                     "payload": {
#                         "num": chunk_num,
#                         "data": b64_data
#                     }
#                 }
#                 # 序列化JSON（自动处理特殊字符）
#                 try:
#                     frame_json = json.dumps(data_frame, ensure_ascii=False)
#                 except Exception as e:
#                     print(f"数据帧JSON序列化失败：{e}")
#                     return
#                 # 发送前检查JSON格式（可选，用于调试）
#                 # print(f"发送数据帧 {chunk_num}：{frame_json[:500]}...")
#                 ws.send(frame_json)
#                 print(f"发送数据帧 {chunk_num}/{total_chunks}", end='\r')
#
#         # 3. 发送结束帧（如果需要补充数据）
#         # 注：本实现中数据帧已包含完整数据，结束帧仅作标记
#         end_frame = {
#             "header": {
#                 "user_id": "default_user",
#                 "uuid": file_uuid,
#                 "status": 2
#             },
#             "payload": {
#                 "num": total_chunks,
#                 "size": str(file_size),
#                 "data": ""
#             }
#         }
#         ws.send(json.dumps(end_frame))
#         print("已发送结束帧")
#
#     def on_close(ws, close_status_code, close_msg):
#         print(f"连接关闭: {close_status_code}, {close_msg}")
#
#     def on_error(ws, error):
#         print(f"错误: {error}")
#
#     ws = websocket.WebSocketApp(
#         server_url,
#         on_open=on_open,
#         on_message=on_message,
#         on_close=on_close,
#         on_error=on_error
#     )
#     ws.run_forever()
#
# if __name__ == "__main__":
#     send_file_with_protocol("./data/sleep_data_2025年09月16日 22.00-2025年09月17日 09.00.txt")

import websocket
import json
import os
import hashlib
import uuid
import base64
import glob
import time
WS_IP_ADDRESS_PORT = "ws://47.106.145.145:50232"
current_file = None

def calculate_md5(file_path):
    md5_hash = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            md5_hash.update(chunk)
    return md5_hash.hexdigest()

def on_message(ws, message):
    print(f"[testPspSleepV3Demo] 收到原始消息: {message[:200]}...")  # 打印前200字符
    global current_file
    try:
        frame = json.loads(message)
        header = frame.get("header", {})
        payload = frame.get("payload", {})
        status = header.get("status")

        print(f"[testPspSleepV3Demo] 收到消息状态: {status}, 文件: {header.get('fileName')}")

        if status == 3:
            current_file = {
                "uuid": header["uuid"],
                "filename": header["fileName"],
                "total_size": int(header["size"]),
                "md5": header["MD5"],
                "total_chunks": payload["packageNum"],
                "chunks": {}
            }
            print(f"[testPspSleepV3Demo] 元数据帧初始化: {current_file}")
        elif status == 4:
            if not current_file:
                print("[testPspSleepV3Demo] 收到数据帧但未初始化 current_file")
                return
            chunk_num = payload["num"]
            data = base64.b64decode(payload["data"])
            current_file["chunks"][chunk_num] = data
            print(f"[testPspSleepV3Demo] 收到数据块 {chunk_num}/{current_file['total_chunks']}")
        elif status == 5:
            if not current_file:
                print("[testPspSleepV3Demo] 收到结束帧但未初始化 current_file")
                return
            save_dir = "./received_from_server"
            os.makedirs(save_dir, exist_ok=True)
            file_path = os.path.join(save_dir, current_file["filename"])
            with open(file_path, "wb") as f:
                for i in range(current_file["total_chunks"]):
                    if i not in current_file["chunks"]:
                        print(f"[testPspSleepV3Demo] 缺失数据块 {i}")
                        return
                    f.write(current_file["chunks"][i])
            actual_md5 = calculate_md5(file_path)
            if actual_md5 != current_file["md5"]:
                print(f"[testPspSleepV3Demo] MD5 校验失败: {actual_md5} != {current_file['md5']}")
                os.remove(file_path)
            else:
                print(f"[testPspSleepV3Demo] 文件接收完成并校验通过: {file_path}")
            current_file = None
    except Exception as e:
        print(f"处理消息出错: {e}")

def send_file_with_protocol(file_path, server_url=WS_IP_ADDRESS_PORT):
    if not os.path.isfile(file_path):
        print(f"[testPspSleepV3Demo] 文件不存在: {file_path}")
        return None

    file_size = os.path.getsize(file_path)
    file_md5 = calculate_md5(file_path)
    file_uuid = str(uuid.uuid4())
    chunk_size = 1024 * 32
    total_chunks = (file_size + chunk_size - 1) // chunk_size

    print(f"[testPspSleepV3Demo] 准备发送文件: {file_path}, 大小: {file_size} bytes, 总块数: {total_chunks}")

    def on_open(ws):
        print("[testPspSleepV3Demo] WebSocket 已连接，开始发送文件")
        first_frame = {
            "header": {
                "user_id": "default_user",
                "uuid": file_uuid,
                "msgType": "txt",
                "fileName": file_path.rsplit('/', 1)[-1],
                "size": str(file_size),
                "MD5": file_md5,
                "status": 0
            },
            "payload": {"packageNum": total_chunks, "data": ""}
        }
        ws.send(json.dumps(first_frame))
        print("[testPspSleepV3Demo] 已发送元数据帧")

        with open(file_path, "rb") as f:
            for chunk_num in range(total_chunks):
                data = f.read(chunk_size)
                # b64_data = base64.b64encode(data).decode('utf-8').replace('\n', '').replace('\r', '')
                b64_data = base64.b64encode(data).decode('utf-8')
                data_frame = {
                    "header": {"user_id": "default_user", "uuid": file_uuid, "status": 1},
                    "payload": {"num": chunk_num, "data": b64_data}
                }
                ws.send(json.dumps(data_frame))
                print(f"发送数据帧 {chunk_num}/{total_chunks}", end='\r')
                # if chunk_num % 10 == 0:
                #     print(f"[testPspSleepV3Demo] 已发送数据块 {chunk_num}/{total_chunks}")

        end_frame = {
            "header": {"user_id": "default_user", "uuid": file_uuid, "status": 2},
            "payload": {"num": total_chunks, "size": str(file_size), "data": ""}
        }
        ws.send(json.dumps(end_frame))
        print("[testPspSleepV3Demo] 已发送结束帧")
    def on_close(ws, close_status_code, close_msg):
        print(f"连接关闭: {close_status_code}, {close_msg}")

    def on_error(ws, error):
        print(f"错误: {error}")

    ws = websocket.WebSocketApp(server_url, on_open=on_open, on_message=on_message,on_close=on_close,
        on_error=on_error)
    ws.run_forever()

    files = glob.glob("./received_from_server/*.txt")
    if files:
        latest = sorted(files, key=os.path.getmtime)[-1]
        print(f"[testPspSleepV3Demo] 返回最新接收文件: {latest}")
        return latest
    return None

def run_upload(file_path):
    print(f"[run_upload] 开始上传: {file_path}")
    # for f in glob.glob("./received_from_server/*.txt"):
    #     os.remove(f)
    #     print(f"[run_upload] 删除旧文件: {f}")
    returned_txt = send_file_with_protocol(file_path)
    if returned_txt:
        print(f"[run_upload] 上传完成，返回文件: {returned_txt}")
    else:
        print("[run_upload] 上传失败或未返回文件")
    return returned_txt
